package com.cx.core.redis.service;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Service;

import java.util.Collections;

@Service
public class LimitService {

    @Autowired
    private RedisTemplate redisTemplate;

    public boolean limit(String key, int limit, int time) {
        String luaScript =  "local key = KEYS[1] -- 限流资源\n" +
                "local maxRequestCount = ARGV[1] -- 限流请求数\n" +
                "local windowSize = ARGV[2] -- 限流时间\n" +
                "local currentCount = redis.call('get', key) -- 当前请求数\n" +
                "-- 限流存在并且超过限流大小，则返回剩余可用请求数=0\n" +
                "if (currentCount and tonumber(currentCount) >= tonumber(maxRequestCount)) then\n" +
                "    return -1\n" +
                "end\n" +
                "-- 请求数自增\n" +
                "currentCount = redis.call('incr', key)\n" +
                "-- 第一次请求，则设置过期时间\n" +
                "if (tonumber(currentCount) == 1) then\n" +
                "    redis.call('expire', key, windowSize)\n" +
                "end\n" +
                "-- 返回剩余可用请求数\n" +
                "return tonumber(maxRequestCount) - tonumber(currentCount)\n";
        RedisScript redisScript = RedisScript.of(luaScript);
        Number result = (Number) redisTemplate.execute(redisScript, Collections.singletonList(key), limit, time);
        return result.intValue() > 0;
    }
}
